#include "message_decode_buffer.h"
#include <assert.h>
#include <string>
#include <utility>
#include <set>

#include "fecEncoderRS8.h"
#include "fecEncoderRS16.h"

class BlockDecodeBuffer
{
public:
    BlockDecodeBuffer(uint64_t block_id, int block_size, int parity_size);
    void reset(uint64_t block_id, int block_size, int parity_size);
    bool ready();
    MessageList appendAndTryDecode(const char* buf, int len);

private:
    uint64_t    block_id_;
    int         block_size_;
    int         parity_size_;
    //Message     msgs_;
    FecPacketList     loss_packets_;
    std::map<int, SpFecPacket> msgs_;
};

BlockDecodeBuffer::BlockDecodeBuffer(uint64_t block_id, int block_size, int parity_size)
    :   block_id_(block_id),
        block_size_(block_size),
        parity_size_(parity_size)
{ }

void BlockDecodeBuffer::reset(uint64_t block_id, int block_size, int parity_size)
{
    block_id_ = block_id;
    block_size_ = block_size;
    parity_size_ = parity_size;
    msgs_.clear();
    loss_packets_.clear(); 
}

bool BlockDecodeBuffer::ready() { return msgs_.size() >= static_cast<size_t>(block_size_); }

MessageList BlockDecodeBuffer::appendAndTryDecode(const char* buf, int len)
{
    SpFecPacket sppacket = std::make_shared<FecPacket>();
    memcpy(&(*sppacket), buf, len);

    //LOG_PRINT(MAIN, info, "fums BlockDecodeBuffer::appendAndTryDecode [{}:{}:{}]", 
    //    block_id_, sppacket->segment_id(), sppacket->get_msg_id());

    assert(block_id_ == sppacket->block_id());

    std::map<int, SpFecPacket>::iterator iter = msgs_.find(sppacket->segment_id());
    if (iter != msgs_.end())
    {
        SpFecPacket p = iter->second;
        //LOG_PRINT(MAIN, warning, "fums BlockDecodeBuffer::appendAndTryDecode recv duplication seg [{}:{}:{}], {}:{}, {}", 
        //    block_id_, sppacket->segment_id(), sppacket->get_msg_id(), p->get_packet_length(), sppacket->get_packet_length(), 
        //    memcmp(p->get_packet_payload(), sppacket->get_packet_payload(), p->get_packet_length()) == 0);

        assert(p->get_packet_length() == sppacket->get_packet_length());
        assert(memcmp(p->get_packet_payload(), sppacket->get_packet_payload(), p->get_packet_length()) == 0);

        return MessageList();
    }
    msgs_[sppacket->segment_id()] = sppacket;
                      
    if (!ready())
        return MessageList();

    std::unique_ptr<uint8_t *[]> rxPtr(new uint8_t *[block_size_ + parity_size_]);
    memset(&(rxPtr.get()[0]), 0, sizeof(uint8_t *) * (block_size_ + parity_size_));

    for (std::map<int, SpFecPacket>::iterator iter = msgs_.begin(); iter != msgs_.end(); ++iter)
    {
        int seg_id = iter->second->segment_id();
        rxPtr.get()[seg_id] = iter->second->get_fecpayload();

    }

    std::unique_ptr<unsigned int[]> erasureLocs(new unsigned int[parity_size_]);
    int erasureLocs_index = 0;
    for (int i = 0; i < block_size_ + parity_size_; ++i)
    {
        if (rxPtr.get()[i] == NULL)
        {
            erasureLocs.get()[erasureLocs_index++] = i;
            SpFecPacket new_packet = std::make_shared<FecPacket>();
            rxPtr.get()[i] = new_packet->get_fecpayload();
            loss_packets_.push_back(new_packet);
        }
    }
    assert(loss_packets_.size() == parity_size_);

    Fec::FecDecoderRS16 decoder;
    int ret = decoder.Init(block_size_, parity_size_, sizeof(UdpFecPacketMsg));
    assert(ret);
    ret = decoder.Decode((char**)rxPtr.get(), block_size_, parity_size_, erasureLocs.get());
    assert(ret != 0);

    MessageList msgs;
    int current_msg_id = 0;
    std::string new_msg;

    int start_seg_index = 0;
    for (int i = 0; i < block_size_; ++i)
    {
        UdpFecPacketMsg *msg = (UdpFecPacketMsg *)(rxPtr.get()[i]);
        if (current_msg_id == msg->get_msg_id())
        {
            new_msg.append((char*)msg->get_msg_data(), msg->get_msg_data_length());
        }
        else
        {
            //LOG_PRINT(MAIN, info, "fums BlockDecodeBuffer::appendAndTryDecode [{}:{}-{}] msg complete",
            //    block_id_, start_seg_index, i);
            current_msg_id++;
            start_seg_index = i;
            msgs.push_back(std::move(new_msg));
            new_msg.clear();
            new_msg.append((char*)msg->get_msg_data(), msg->get_msg_data_length());
        }
    }
    //LOG_PRINT(MAIN, info, "fums BlockDecodeBuffer::appendAndTryDecode [{}:{}-{}] msg complete",
    //    block_id_, start_seg_index, block_size_);

    assert(!new_msg.empty());
    msgs.push_back(std::move(new_msg));

    return msgs;
}

MessageDecodeBuffer::MessageDecodeBuffer(int window_size) :
    window_size_(window_size),
    next_block_id_(0)
{ 
    for (size_t i = 0; i < window_size_; ++i)
    {
        free_buffs_.push_back(std::make_shared<BlockDecodeBuffer>(0, 0, 0));
    }
}

MessageDecodeBuffer::~MessageDecodeBuffer(){}
    
SpBlockDecodeBuffer MessageDecodeBuffer::addNewBlock(FecPacket* packet)
{
    SpBlockDecodeBuffer new_block;
    if (full())
    {
        std::map<uint64_t, SpBlockDecodeBuffer>::iterator iter = buffs_.begin();
        //LOG_PRINT(MAIN, info, "fums BlockDecodeBuffer::addNewBlock drop block {}", iter->first);
        free_buffs_.push_back(iter->second);
        buffs_.erase(iter);
    }
    
    assert(!free_buffs_.empty());
    new_block = free_buffs_.front();
    free_buffs_.pop_front();
    new_block->reset(packet->block_id(), packet->block_size(), packet->parity_num());
    //LOG_PRINT(MAIN, info, "fums BlockDecodeBuffer::addNewBlock add block {}", packet->block_id());
    buffs_[packet->block_id()] = new_block;
    return new_block;
}

MessageList MessageDecodeBuffer::appendAndTryDecode(const char* buf, int len)
{
    assert(NULL != buf && len > 0);

    MessageList  packets;
    FecPacket* packet = (FecPacket*)buf;

    //LOG_PRINT(MAIN, info, "fums MessageDecodeBuffer::appendAndTryDecode [{}:{}:{}]", packet->block_id(), packet->segment_id(),
    //    packet->get_msg_id());

    if (packet->block_id() < getMinBlockId())
        return packets;

    SpBlockDecodeBuffer new_buff = find(packet->block_id());
    if (!new_buff)
    {
        new_buff = addNewBlock(packet);
    }

    packets = new_buff->appendAndTryDecode(buf, len);
    if (packets.size() > 0)
    {
        next_block_id_ = packet->block_id() + 1;
        freeBlockBuffer(packet->block_id());
    }

    return packets;
}


uint64_t MessageDecodeBuffer::getMaxBlockId() 
{
    std::map<uint64_t, SpBlockDecodeBuffer>::reverse_iterator iter = buffs_.rbegin();
    return iter != buffs_.rend() ? iter->first : 0;
}

uint64_t MessageDecodeBuffer::getMinBlockId() 
{
    return next_block_id_;
}

int MessageDecodeBuffer::full()
{
    assert(buffs_.size() <= window_size_);
    return buffs_.size() == window_size_;
}

SpBlockDecodeBuffer MessageDecodeBuffer::find(uint64_t blockid)
{
    std::map<uint64_t, SpBlockDecodeBuffer>::iterator iter = buffs_.find(blockid);
    return iter == buffs_.end() ? SpBlockDecodeBuffer() : iter->second;
}

void MessageDecodeBuffer::freeBlockBuffer(uint64_t  blockid)
{
    std::map<uint64_t, SpBlockDecodeBuffer>::iterator iter = buffs_.find(blockid);
    assert(iter != buffs_.end());
    free_buffs_.push_back(iter->second);
    buffs_.erase(iter);
}